from PyQt6.QtWidgets import (
    QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, 
    QPushButton, QTreeWidget, QTreeWidgetItem, QTabWidget, 
    QLabel, QSplitter, QGroupBox, QFormLayout, QTableWidget, 
    QTableWidgetItem, QHeaderView, QMessageBox, QMenu, QToolButton,
    QLineEdit, QCheckBox, QInputDialog
)
from PyQt6.QtCore import Qt, QSize, pyqtSignal, QTimer
from PyQt6.QtGui import QIcon, QAction, QCursor

import os
import sys
from typing import Dict, List, Any, Optional

from ui.connection_dialog import ConnectionDialog
from ui.collection_tab import CollectionTab
from utils.config_manager import ConfigManager
from core.chroma_client import ChromaClient


class ConnectionTreeItem(QTreeWidgetItem):
    """连接树项目"""
    
    def __init__(self, connection: Dict[str, Any], parent=None):
        """
        初始化连接树项目
        
        Args:
            connection: 连接配置
            parent: 父项目
        """
        super().__init__(parent)
        self.connection = connection
        self.setText(0, connection.get('name', '未命名连接'))
        self.setToolTip(0, f"{connection.get('type', 'remote')} - {connection.get('host', '')}")


class CollectionTreeItem(QTreeWidgetItem):
    """集合树项目"""
    
    def __init__(self, name: str, parent=None):
        """
        初始化集合树项目
        
        Args:
            name: 集合名称
            parent: 父项目
        """
        super().__init__(parent)
        self.collection_name = name
        self.setText(0, name)


class ConnectionTab(QWidget):
    """连接标签页"""
    
    # 自定义信号，用于关闭标签页
    closeRequested = pyqtSignal(str)
    # 自定义信号，用于打开集合
    openCollectionRequested = pyqtSignal(str, object)
    
    def __init__(self, connection: Dict[str, Any], parent=None):
        """
        初始化连接标签页
        
        Args:
            connection: 连接配置
            parent: 父窗口
        """
        super().__init__(parent)
        self.connection = connection
        self.client = None
        self.collections = []
        
        self._init_ui()
        self._connect_to_server()
    
    def _init_ui(self):
        """初始化UI组件"""
        main_layout = QVBoxLayout()
        main_layout.setContentsMargins(10, 10, 10, 10)
        
        # 创建上部布局
        top_layout = QHBoxLayout()
        
        # 自动刷新选项
        self.auto_refresh_check = QCheckBox("自动刷新")
        self.auto_refresh_check.setChecked(True)
        self.auto_refresh_check.stateChanged.connect(self._on_auto_refresh_changed)
        top_layout.addWidget(self.auto_refresh_check)
        
        top_layout.addStretch()
        
        # 添加到主布局
        main_layout.addLayout(top_layout)
        
        # 创建信息面板
        info_layout = QHBoxLayout()
        
        # 服务器信息面板
        self.server_group = QGroupBox("服务器")
        server_layout = QFormLayout()
        
        self.server_version_label = QLabel("ChromaDB版本: 未知")
        server_layout.addRow(self.server_version_label)
        
        self.server_os_label = QLabel("OS: 未知")
        server_layout.addRow(self.server_os_label)
        
        self.server_process_label = QLabel("进程ID: 未知")
        server_layout.addRow(self.server_process_label)
        
        self.server_group.setLayout(server_layout)
        
        # 状态信息面板
        self.status_group = QGroupBox("状态")
        status_layout = QFormLayout()
        
        self.status_clients_label = QLabel("连接状态: 已连接")
        status_layout.addRow(self.status_clients_label)
        
        self.status_group.setLayout(status_layout)
        
        # 添加到信息布局
        info_layout.addWidget(self.server_group)
        info_layout.addWidget(self.status_group)
        
        # 添加到主布局
        main_layout.addLayout(info_layout)
        
        # 集合统计面板
        self.collections_group = QGroupBox("集合统计")
        collections_layout = QVBoxLayout()
        
        # 集合操作按钮
        collection_btn_layout = QHBoxLayout()
        
        self.create_collection_btn = QPushButton("创建集合")
        self.create_collection_btn.clicked.connect(self._on_create_collection)
        collection_btn_layout.addWidget(self.create_collection_btn)
        
        self.delete_collection_btn = QPushButton("删除集合")
        self.delete_collection_btn.clicked.connect(self._on_delete_collection)
        collection_btn_layout.addWidget(self.delete_collection_btn)
        
        collections_layout.addLayout(collection_btn_layout)
        
        self.collections_table = QTableWidget()
        self.collections_table.setColumnCount(4)
        self.collections_table.setHorizontalHeaderLabels(["集合名称", "文档数", "嵌入模型", "创建时间"])
        self.collections_table.horizontalHeader().setSectionResizeMode(QHeaderView.ResizeMode.Stretch)
        self.collections_table.setSelectionBehavior(QTableWidget.SelectionBehavior.SelectRows)
        self.collections_table.setSelectionMode(QTableWidget.SelectionMode.SingleSelection)
        self.collections_table.doubleClicked.connect(self._on_collection_double_clicked)
        
        collections_layout.addWidget(self.collections_table)
        self.collections_group.setLayout(collections_layout)
        
        # 添加到主布局
        main_layout.addWidget(self.collections_group)
        
        self.setLayout(main_layout)
        
        # 创建定时器用于自动刷新
        self.refresh_timer = QTimer(self)
        self.refresh_timer.timeout.connect(self.refresh_data)
        self.refresh_timer.start(5000)  # 5秒刷新一次
    
    def _connect_to_server(self):
        """连接到ChromaDB服务器"""
        try:
            self.client = ChromaClient(self.connection)
            self.refresh_data()
        except Exception as e:
            QMessageBox.critical(self, "连接错误", f"无法连接到ChromaDB服务器： {str(e)}")
            self.closeRequested.emit(self.connection.get('name', ''))
    
    def refresh_data(self):
        """刷新数据"""
        if not self.client:
            return
        
        try:
            # 获取集合列表
            collections = self.client.list_collections()
            self.collections = collections
            
            # 更新集合表格
            self.collections_table.setRowCount(len(collections))
            for i, collection in enumerate(collections):
                # 集合名称
                name_item = QTableWidgetItem(collection.name)
                self.collections_table.setItem(i, 0, name_item)
                
                # 文档数
                count = collection.count()
                count_item = QTableWidgetItem(str(count))
                self.collections_table.setItem(i, 1, count_item)
                
                # 嵌入模型
                model = collection.metadata.get('embedding_model', '未知')
                model_item = QTableWidgetItem(model)
                self.collections_table.setItem(i, 2, model_item)
                
                # 创建时间 (ChromaDB目前不提供创建时间，使用占位符)
                time_item = QTableWidgetItem("未知")
                self.collections_table.setItem(i, 3, time_item)
            
            # 更新服务器信息
            self.server_version_label.setText(f"ChromaDB类型: {self.client.client.__class__.__name__}")
            self.server_os_label.setText(f"系统: {os.name}")
            self.server_process_label.setText(f"进程ID: {os.getpid()}")
            
            # 更新状态信息
            self.status_clients_label.setText("连接状态: 已连接")
            
        except Exception as e:
            QMessageBox.warning(self, "刷新错误", f"刷新数据时出错: {str(e)}")
    
    def _on_auto_refresh_changed(self, state):
        """
        自动刷新状态变更处理
        
        Args:
            state: 复选框状态
        """
        if state == Qt.CheckState.Checked:
            self.refresh_timer.start(5000)
        else:
            self.refresh_timer.stop()
    
    def _on_create_collection(self):
        """创建集合按钮点击处理"""
        if not self.client:
            QMessageBox.warning(self, "错误", "未连接到服务器")
            return
        
        name, ok = QInputDialog.getText(self, "创建集合", "集合名称:")
        if ok and name:
            try:
                self.client.create_collection(name)
                self.refresh_data()
                QMessageBox.information(self, "成功", f"集合 '{name}' 创建成功")
            except Exception as e:
                QMessageBox.critical(self, "创建错误", f"创建集合时出错: {str(e)}")
    
    def _on_delete_collection(self):
        """删除集合按钮点击处理"""
        if not self.client:
            QMessageBox.warning(self, "错误", "未连接到服务器")
            return
        
        # 获取选中的行
        selected_rows = self.collections_table.selectionModel().selectedRows()
        if not selected_rows:
            QMessageBox.warning(self, "选择错误", "请选择要删除的集合")
            return
        
        row = selected_rows[0].row()
        name = self.collections_table.item(row, 0).text()
        
        # 确认对话框
        reply = QMessageBox.question(
            self, 
            "确认删除", 
            f"确定要删除集合 '{name}' 吗？此操作不可撤销！",
            QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No,
            QMessageBox.StandardButton.No
        )
        
        if reply == QMessageBox.StandardButton.Yes:
            try:
                self.client.delete_collection(name)
                self.refresh_data()
                QMessageBox.information(self, "成功", f"集合 '{name}' 删除成功")
                
                # 发送信号通知主窗口关闭相关标签页
                self.openCollectionRequested.emit(name, None)
            except Exception as e:
                QMessageBox.critical(self, "删除错误", f"删除集合时出错: {str(e)}")
    
    def _on_collection_double_clicked(self, index):
        """
        集合双击处理
        
        Args:
            index: 表格索引
        """
        if not self.client:
            return
        
        row = index.row()
        name = self.collections_table.item(row, 0).text()
        
        # 发送信号通知主窗口打开集合标签页
        self.openCollectionRequested.emit(name, self.client)


class MainWindow(QMainWindow):
    """ChromaDB管理器主窗口"""
    
    def __init__(self):
        """初始化主窗口"""
        super().__init__()
        
        self.config_manager = ConfigManager()
        self.connections = {}  # 存储已打开的连接
        self.collection_tabs = {}  # 存储已打开的集合标签页
        
        self._init_ui()
        self._load_connections()
    
    def _init_ui(self):
        """初始化UI组件"""
        self.setWindowTitle("ChromaDB Desktop Manager")
        self.resize(1200, 800)
        
        # 创建中央部件
        central_widget = QWidget()
        main_layout = QHBoxLayout()
        
        # 创建分割器
        splitter = QSplitter(Qt.Orientation.Horizontal)
        
        # 左侧面板
        left_panel = QWidget()
        left_layout = QVBoxLayout()
        
        # 新建连接按钮
        self.new_connection_btn = QPushButton("新建连接")
        self.new_connection_btn.setIcon(QIcon.fromTheme("list-add"))
        self.new_connection_btn.clicked.connect(self._on_new_connection)
        left_layout.addWidget(self.new_connection_btn)
        
        # 连接树
        self.connection_tree = QTreeWidget()
        self.connection_tree.setHeaderHidden(True)
        self.connection_tree.setContextMenuPolicy(Qt.ContextMenuPolicy.CustomContextMenu)
        self.connection_tree.customContextMenuRequested.connect(self._on_tree_context_menu)
        self.connection_tree.itemDoubleClicked.connect(self._on_connection_double_clicked)
        left_layout.addWidget(self.connection_tree)
        
        left_panel.setLayout(left_layout)
        
        # 右侧面板
        right_panel = QWidget()
        right_layout = QVBoxLayout()
        
        # 标签页
        self.tab_widget = QTabWidget()
        self.tab_widget.setTabsClosable(True)
        self.tab_widget.tabCloseRequested.connect(self._on_tab_close)
        right_layout.addWidget(self.tab_widget)
        
        right_panel.setLayout(right_layout)
        
        # 添加到分割器
        splitter.addWidget(left_panel)
        splitter.addWidget(right_panel)
        splitter.setSizes([200, 1000])  # 设置初始大小
        
        # 添加到主布局
        main_layout.addWidget(splitter)
        central_widget.setLayout(main_layout)
        
        self.setCentralWidget(central_widget)
    
    def _load_connections(self):
        """加载连接配置"""
        connections = self.config_manager.get_connections()
        
        # 清空连接树
        self.connection_tree.clear()
        
        # 添加连接到树
        for connection in connections:
            item = ConnectionTreeItem(connection)
            self.connection_tree.addTopLevelItem(item)
    
    def _on_new_connection(self):
        """新建连接按钮点击处理"""
        dialog = ConnectionDialog(self)
        dialog.connectionAccepted.connect(self._on_connection_accepted)
        dialog.exec()
    
    def _on_connection_accepted(self, connection: Dict[str, Any]):
        """
        连接配置接受处理
        
        Args:
            connection: 连接配置
        """
        # 添加到配置管理器
        self.config_manager.add_connection(connection)
        
        # 添加到连接树
        item = ConnectionTreeItem(connection)
        self.connection_tree.addTopLevelItem(item)
        
        # 打开连接
        self._open_connection(connection)
    
    def _on_connection_double_clicked(self, item):
        """
        连接项目双击处理
        
        Args:
            item: 树项目
        """
        if isinstance(item, ConnectionTreeItem):
            self._open_connection(item.connection)
    
    def _open_connection(self, connection: Dict[str, Any]):
        """
        打开连接
        
        Args:
            connection: 连接配置
        """
        name = connection.get('name', '')
        
        # 检查是否已打开
        for i in range(self.tab_widget.count()):
            if self.tab_widget.tabText(i) == name:
                self.tab_widget.setCurrentIndex(i)
                return
        
        try:
            # 创建新标签页
            tab = ConnectionTab(connection)
            tab.closeRequested.connect(self._on_tab_close_by_name)
            tab.openCollectionRequested.connect(self._on_open_collection)
            
            # 添加到标签页
            index = self.tab_widget.addTab(tab, name)
            self.tab_widget.setCurrentIndex(index)
            
            # 存储连接
            self.connections[name] = tab
        except Exception as e:
            QMessageBox.critical(self, "连接错误", f"打开连接时出错: {str(e)}")
    
    def _on_open_collection(self, collection_name: str, client: Optional[ChromaClient]):
        """
        打开集合
        
        Args:
            collection_name: 集合名称
            client: ChromaDB客户端
        """
        if not client:
            # 如果client为None，则关闭集合标签页
            self._close_collection_tab(collection_name)
            return
        
        # 生成标签页ID
        tab_id = f"collection:{collection_name}"
        
        # 检查是否已打开
        for i in range(self.tab_widget.count()):
            if self.tab_widget.tabText(i) == collection_name:
                self.tab_widget.setCurrentIndex(i)
                return
        
        # 创建新标签页
        tab = CollectionTab(client, collection_name)
        
        # 添加到标签页
        index = self.tab_widget.addTab(tab, collection_name)
        self.tab_widget.setCurrentIndex(index)
        
        # 存储集合标签页
        self.collection_tabs[tab_id] = tab
    
    def _close_collection_tab(self, collection_name: str):
        """
        关闭集合标签页
        
        Args:
            collection_name: 集合名称
        """
        tab_id = f"collection:{collection_name}"
        
        # 检查是否已打开
        for i in range(self.tab_widget.count()):
            if self.tab_widget.tabText(i) == collection_name:
                self.tab_widget.removeTab(i)
                
                # 从存储中移除
                if tab_id in self.collection_tabs:
                    del self.collection_tabs[tab_id]
                break
    
    def _on_tab_close(self, index: int):
        """
        标签页关闭处理
        
        Args:
            index: 标签页索引
        """
        name = self.tab_widget.tabText(index)
        self.tab_widget.removeTab(index)
        
        # 从存储中移除
        if name in self.connections:
            del self.connections[name]
    
    def _on_tab_close_by_name(self, name: str):
        """
        通过名称关闭标签页
        
        Args:
            name: 标签页名称
        """
        for i in range(self.tab_widget.count()):
            if self.tab_widget.tabText(i) == name:
                self.tab_widget.removeTab(i)
                
                # 从存储中移除
                if name in self.connections:
                    del self.connections[name]
                break
    
    def _on_tree_context_menu(self, pos):
        """
        树上下文菜单处理
        
        Args:
            pos: 菜单位置
        """
        item = self.connection_tree.itemAt(pos)
        if not item:
            return
        
        menu = QMenu(self)
        
        if isinstance(item, ConnectionTreeItem):
            # 连接项目菜单
            open_action = QAction("打开连接", self)
            open_action.triggered.connect(lambda: self._open_connection(item.connection))
            menu.addAction(open_action)
            
            edit_action = QAction("编辑连接", self)
            edit_action.triggered.connect(lambda: self._edit_connection(item))
            menu.addAction(edit_action)
            
            delete_action = QAction("删除连接", self)
            delete_action.triggered.connect(lambda: self._delete_connection(item))
            menu.addAction(delete_action)
        
        menu.exec(QCursor.pos())
    
    def _edit_connection(self, item: ConnectionTreeItem):
        """
        编辑连接
        
        Args:
            item: 连接树项目
        """
        dialog = ConnectionDialog(self, item.connection)
        
        def on_accepted(connection):
            # 更新配置管理器
            self.config_manager.update_connection(item.connection.get('name', ''), connection)
            
            # 更新树项目
            item.connection = connection
            item.setText(0, connection.get('name', '未命名连接'))
            item.setToolTip(0, f"{connection.get('type', 'remote')} - {connection.get('host', '')}")
            
            # 关闭并重新打开标签页
            old_name = item.connection.get('name', '')
            new_name = connection.get('name', '')
            
            if old_name in self.connections:
                self._on_tab_close_by_name(old_name)
                self._open_connection(connection)
        
        dialog.connectionAccepted.connect(on_accepted)
        dialog.exec()
    
    def _delete_connection(self, item: ConnectionTreeItem):
        """
        删除连接
        
        Args:
            item: 连接树项目
        """
        name = item.connection.get('name', '')
        
        # 确认对话框
        reply = QMessageBox.question(
            self, 
            "确认删除", 
            f"确定要删除连接 '{name}' 吗？",
            QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No,
            QMessageBox.StandardButton.No
        )
        
        if reply == QMessageBox.StandardButton.Yes:
            # 从配置管理器中删除
            self.config_manager.delete_connection(name)
            
            # 从树中删除
            index = self.connection_tree.indexOfTopLevelItem(item)
            self.connection_tree.takeTopLevelItem(index)
            
            # 关闭标签页
            self._on_tab_close_by_name(name)
